# -*- coding: utf-8 -*-
"""
Locust压力测试文件

使用Locust框架对API进行压力测试，模拟真实用户行为
支持多种测试场景：基础功能测试、用户管理操作、混合场景测试

运行命令:
    locust -f tests/performance/locustfile.py --host=http://localhost:8000
    locust -f tests/performance/locustfile.py --host=http://localhost:8000 --users=100 --spawn-rate=10 --run-time=60s
"""

import json
import random
import time
from locust import HttpUser, task, between, events
from locust.exception import RescheduleTask
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class APIUser(HttpUser):
    """
    API用户行为类
    
    模拟真实用户对API的访问行为，包括：
    - 基础接口访问
    - 用户管理操作
    - 错误处理测试
    """
    
    # 用户请求间隔时间（秒）
    wait_time = between(1, 3)
    
    # 测试数据
    test_users = []
    created_user_ids = []
    
    def on_start(self):
        """
        用户开始测试时的初始化操作
        """
        logger.info(f"用户 {self.client.base_url} 开始压力测试")
        
        # 预生成测试用户数据
        for i in range(10):
            user_data = {
                "name": f"test_user_{random.randint(1000, 9999)}",
                "email": f"test{random.randint(1000, 9999)}@example.com",
                "age": random.randint(18, 65)
            }
            self.test_users.append(user_data)
    
    def on_stop(self):
        """
        用户停止测试时的清理操作
        """
        logger.info(f"用户 {self.client.base_url} 结束压力测试")
        
        # 清理创建的测试数据
        for user_id in self.created_user_ids:
            try:
                self.client.delete(f"/api/v1/users/{user_id}")
            except Exception as e:
                logger.warning(f"清理用户数据失败: {e}")
    
    @task(10)
    def test_root_endpoint(self):
        """
        测试根路径接口
        权重: 10 (高频访问)
        """
        with self.client.get("/", catch_response=True) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"根路径访问失败: {response.status_code}")
    
    @task(8)
    def test_health_check(self):
        """
        测试健康检查接口
        权重: 8 (高频访问)
        """
        with self.client.get("/health", catch_response=True) as response:
            if response.status_code == 200:
                try:
                    data = response.json()
                    if data.get("status") == "healthy":
                        response.success()
                    else:
                        response.failure("健康检查状态异常")
                except json.JSONDecodeError:
                    response.failure("健康检查响应格式错误")
            else:
                response.failure(f"健康检查失败: {response.status_code}")
    
    @task(5)
    def test_get_all_users(self):
        """
        测试获取所有用户接口
        权重: 5 (中频访问)
        """
        with self.client.get("/api/v1/users", catch_response=True) as response:
            if response.status_code == 200:
                try:
                    users = response.json()
                    if isinstance(users, list):
                        response.success()
                    else:
                        response.failure("用户列表格式错误")
                except json.JSONDecodeError:
                    response.failure("用户列表响应格式错误")
            else:
                response.failure(f"获取用户列表失败: {response.status_code}")
    
    @task(3)
    def test_create_user(self):
        """
        测试创建用户接口
        权重: 3 (低频访问)
        """
        if not self.test_users:
            raise RescheduleTask()
        
        user_data = random.choice(self.test_users)
        
        with self.client.post("/api/v1/users", json=user_data, catch_response=True) as response:
            if response.status_code == 201:
                try:
                    created_user = response.json()
                    user_id = created_user.get("id")
                    if user_id:
                        self.created_user_ids.append(user_id)
                        response.success()
                    else:
                        response.failure("创建用户响应缺少ID")
                except json.JSONDecodeError:
                    response.failure("创建用户响应格式错误")
            else:
                response.failure(f"创建用户失败: {response.status_code}")
    
    @task(4)
    def test_get_user_by_id(self):
        """
        测试根据ID获取用户接口
        权重: 4 (中频访问)
        """
        if not self.created_user_ids:
            # 如果没有创建的用户，使用已知存在的用户ID（从数据库中获取）
            user_id = random.randint(1, 10)  # 假设数据库中有ID 1-10的用户
        else:
            user_id = random.choice(self.created_user_ids)
        
        with self.client.get(f"/api/v1/users/{user_id}", catch_response=True) as response:
            if user_id in self.created_user_ids:
                # 测试存在的用户
                if response.status_code == 200:
                    try:
                        user = response.json()
                        if user.get("id") == user_id:
                            response.success()
                        else:
                            response.failure("用户ID不匹配")
                    except json.JSONDecodeError:
                        response.failure("用户详情响应格式错误")
                else:
                    response.failure(f"获取用户详情失败: {response.status_code}")
            else:
                # 测试不存在的用户
                if response.status_code == 404:
                    response.success()
                else:
                    response.failure(f"不存在用户的响应码错误: {response.status_code}")
    
    @task(2)
    def test_update_user(self):
        """
        测试更新用户接口
        权重: 2 (低频访问)
        """
        if not self.created_user_ids:
            raise RescheduleTask()
        
        user_id = random.choice(self.created_user_ids)
        update_data = {
            "name": f"updated_user_{random.randint(1000, 9999)}",
            "email": f"updated{random.randint(1000, 9999)}@example.com",
            "age": random.randint(18, 65)
        }
        
        with self.client.put(f"/api/v1/users/{user_id}", json=update_data, catch_response=True) as response:
            if response.status_code == 200:
                try:
                    updated_user = response.json()
                    if updated_user.get("name") == update_data["name"]:
                        response.success()
                    else:
                        response.failure("用户更新数据不匹配")
                except json.JSONDecodeError:
                    response.failure("更新用户响应格式错误")
            else:
                response.failure(f"更新用户失败: {response.status_code}")
    
    @task(1)
    def test_delete_user(self):
        """
        测试删除用户接口
        权重: 1 (最低频访问)
        """
        if not self.created_user_ids:
            raise RescheduleTask()
        
        user_id = self.created_user_ids.pop()
        
        with self.client.delete(f"/api/v1/users/{user_id}", catch_response=True) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    if result.get("message") == "User deleted successfully":
                        response.success()
                    else:
                        response.failure("删除用户响应消息错误")
                except json.JSONDecodeError:
                    response.failure("删除用户响应格式错误")
            else:
                response.failure(f"删除用户失败: {response.status_code}")
    
    @task(1)
    def test_invalid_endpoints(self):
        """
        测试无效接口
        权重: 1 (测试错误处理)
        """
        invalid_endpoints = [
            "/invalid",
            "/api/v1/users/invalid",
            "/api/v2/users",
            "/nonexistent"
        ]
        
        endpoint = random.choice(invalid_endpoints)
        
        with self.client.get(endpoint, catch_response=True) as response:
            if response.status_code == 404:
                response.success()
            else:
                response.failure(f"无效接口响应码错误: {response.status_code}")


class AdminUser(HttpUser):
    """
    管理员用户行为类
    
    模拟管理员用户的高频操作行为
    """
    
    wait_time = between(0.1, 0.5)
    weight = 1  # 管理员用户权重较低
    
    @task(15)
    def admin_health_check(self):
        """
        管理员频繁检查系统健康状态
        """
        self.client.get("/health")
    
    @task(10)
    def admin_user_management(self):
        """
        管理员进行用户管理操作
        """
        # 获取用户列表
        self.client.get("/api/v1/users")
        
        # 随机查看用户详情
        user_id = random.randint(1, 10)  # 使用已知存在的用户ID范围
        self.client.get(f"/api/v1/users/{user_id}")


class RegularUser(HttpUser):
    """
    普通用户行为类
    
    模拟普通用户的常规访问行为
    """
    
    wait_time = between(2, 5)
    weight = 3  # 普通用户权重较高
    
    @task(20)
    def browse_api(self):
        """
        普通用户浏览API
        """
        self.client.get("/")
    
    @task(5)
    def check_users(self):
        """
        普通用户查看用户信息
        """
        self.client.get("/api/v1/users")


# Locust事件监听器
@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    """
    测试开始时的事件处理
    """
    logger.info("=" * 50)
    logger.info("API压力测试开始")
    logger.info(f"目标主机: {environment.host}")
    logger.info(f"用户数量: {environment.runner.target_user_count if hasattr(environment.runner, 'target_user_count') else '未设置'}")
    
    # 创建测试用户数据
    import requests
    try:
        for i in range(1, 11):  # 创建10个测试用户
            user_data = {
                "name": f"TestUser{i}",
                "email": f"testuser{i}@example.com",
                "age": 20 + i
            }
            response = requests.post(f"{environment.host}/api/v1/users", json=user_data)
            if response.status_code == 201:
                logger.info(f"创建测试用户成功: {user_data['name']}")
            else:
                logger.warning(f"创建测试用户失败: {response.status_code}")
    except Exception as e:
        logger.error(f"创建测试用户时出错: {e}")
    
    logger.info("=" * 50)


@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    """
    测试结束时的事件处理
    """
    logger.info("=" * 50)
    logger.info("API压力测试结束")
    logger.info("测试统计信息:")
    
    stats = environment.stats
    logger.info(f"总请求数: {stats.total.num_requests}")
    logger.info(f"失败请求数: {stats.total.num_failures}")
    logger.info(f"平均响应时间: {stats.total.avg_response_time or 0:.2f}ms")
    logger.info(f"最大响应时间: {stats.total.max_response_time or 0:.2f}ms")
    logger.info(f"最小响应时间: {stats.total.min_response_time or 0:.2f}ms")
    logger.info(f"RPS: {stats.total.current_rps:.2f}")
    logger.info("=" * 50)


@events.request.add_listener
def on_request_failure(request_type, name, response_time, response_length, exception, **kwargs):
    """
    请求失败时的事件处理
    """
    if exception:
        logger.error(f"请求失败: {request_type} {name} - {exception}")


if __name__ == "__main__":
    # 直接运行此文件时的说明
    print("Locust压力测试文件")
    print("请使用以下命令运行压力测试:")
    print("")
    print("基础运行:")
    print("  locust -f tests/performance/locustfile.py --host=http://localhost:8000")
    print("")
    print("指定参数运行:")
    print("  locust -f tests/performance/locustfile.py --host=http://localhost:8000 --users=100 --spawn-rate=10 --run-time=60s")
    print("")
    print("无界面运行:")
    print("  locust -f tests/performance/locustfile.py --host=http://localhost:8000 --users=50 --spawn-rate=5 --run-time=30s --headless")